#  Copyright (c) Meta Platforms, Inc. and affiliates.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
#!/usr/bin/env fbpython
# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.

import torch
from fx2ait.acc_tracer import acc_ops
from fx2ait.tools.common_fx2ait import AITTestCase


class TestPermuteConverter(AITTestCase):
    def test_permute_torch_op(
        self,
    ):
        class TestModule(torch.nn.Module):
            def forward(self, x: torch.Tensor) -> torch.Tensor:
                return torch.permute(x, (2, 0, 1))

        model = TestModule().half().cuda()
        inputs = [torch.randn(32, 256, 256).cuda().half()]
        self.run_test(
            model,
            inputs,
            expected_ops={acc_ops.permute},
        )

    def test_permute_op_on_tensor_tuple(
        self,
    ):
        class TestModule(torch.nn.Module):
            def forward(self, x: torch.Tensor) -> torch.Tensor:
                return x.permute((2, 0, 1))

        model = TestModule().half().cuda()
        inputs = [torch.randn(32, 256, 256).cuda().half()]
        self.run_test(
            model,
            inputs,
            expected_ops={acc_ops.permute},
        )

    def test_permute_op_on_tensor_args(
        self,
    ):
        class TestModule(torch.nn.Module):
            def forward(self, x: torch.Tensor) -> torch.Tensor:
                return x.permute(2, 0, 1)

        model = TestModule().half().cuda()
        inputs = [torch.randn(32, 256, 256).cuda().half()]
        self.run_test(
            model,
            inputs,
            expected_ops={acc_ops.permute},
        )
